#!/usr/bin/env python3
"""Generate V2 subtitle files with smarter line breaks and faster-whisper timing."""
import json, os, re

BASE = r'E:\集群文件夹\factory_os\short_video_real_data_pipeline\phase4b_90s_formal_sample\v22_kimi_claude_loop\full_story_douyin_video'
SUB_DIR = os.path.join(BASE, '04_subtitles', 'v2')

# === CONFIG ===
TOTAL_DURATION = 92.0  # from faster-whisper
# Known keywords to highlight
KEYWORDS = [
    '录屏失败', 'AI 跑偏', '跑偏', '真实执行日志', '真实日志', 'Kimi',
    '69', '72', '81', '85', '89', '88', '92', '76',
    'HTTP 200', '200', '回滚', '冻结', '稳定交付', '自动修复',
    'regression',
]
# Color mapping
KW_COLORS = {
    '录屏失败': '#FF4444', 'AI 跑偏': '#FF8800', '跑偏': '#FF8800',
    '真实执行日志': '#00DDFF', '真实日志': '#00DDFF', 'Kimi': '#00AAFF',
    '69': '#FFD700', '72': '#88FF88', '81': '#44FF44', '85': '#44FFDD',
    '89': '#44DDFF', '88': '#FFDD44', '92': '#FFD700', '76': '#FF2222',
    'HTTP 200': '#44FF44', '200': '#44FF44', '回滚': '#44FF44',
    '冻结': '#44FFDD', '稳定交付': '#44FF44', '自动修复': '#44DDFF',
}

# === SCRIPT: master sentence list with natural grouping ===
# Each entry: (sentence_text, )
SENTENCES = [
    "我让 Claude Code 自己做视频。",
    "录屏失败、AI 跑偏、从 69 爬到 92。",
    "这不是特效，这是真实工作流。",
    "Windows 后台抓不到桌面画面。",
    "试了五种录屏方案，全是空文件。",
    "录屏路线，彻底失败。",
    "更糟的是，Claude Code 跑去研究亚马逊选品。",
    "ChatGPT 看不下去：放弃录屏，让它自己画界面。",
    "于是用 Python 画终端界面、工具调用卡片。",
    "四个 AI 代理并行跑，进度条在涨，代码在敲。",
    "每条命令都来自真实执行日志。",
    "每轮渲染完 Kimi 审核打分。",
    "v1 69 分，改字体 v2 72，加动画 v3 跳到 81。",
    "修字幕 v4 85，砍时长 v6 89。V9 冲到 92。",
    "但下一轮有个硬编码没改。",
    "画面四个代理，底部只显示三个。",
    "92 直接砸到 76。",
    "修了六项回到 88，但回不到 92 了。",
    "一个硬编码，一夜回到解放前。",
    "最终决定：回滚 V9，冻结。",
    "V9 上线审核门户，HTTP 200 通过。",
    "11 轮迭代，1 次崩溃，从 69 到 92。",
    "AI 自己审核自己，自己修复自己，自己部署自己。",
    "录屏失败、跑偏、崩过、又爬回来。",
    "执行、审核、修复、回归、交付。",
    "这才是真正的 AI 工作流。",
]

# === Smart line grouping: combine short sentences (max 2 lines / 13 CJK chars per line) ===
def cjk_chars(s):
    """Count CJK characters in string (excluding ASCII/spaces)."""
    return sum(1 for c in s if '一' <= c <= '鿿')

def total_cjk(s):
    return sum(cjk_chars(c) for c in s)

# Join short sentences into subtitle groups
subtitle_groups = []
current_group = []
current_cjk = 0
current_lines = 0

for sentence in SENTENCES:
    sent_cjk = total_cjk(sentence)

    # If this sentence alone has > 13 CJK chars, split it
    if sent_cjk > 13:
        # Flush current group
        if current_group:
            subtitle_groups.append(list(current_group))
            current_group = []
            current_cjk = 0
            current_lines = 0
        # Split long sentence
        # Try splitting at punctuation
        parts = re.split(r'(?<=[。！？，、])', sentence)
        parts = [p.strip() for p in parts if p.strip()]
        sub_group = []
        sub_cjk = 0
        for part in parts:
            pcjk = total_cjk(part)
            if sub_cjk + pcjk <= 13 and len(sub_group) < 2:
                sub_group.append(part)
                sub_cjk += pcjk
            else:
                if sub_group:
                    subtitle_groups.append(list(sub_group))
                sub_group = [part]
                sub_cjk = pcjk
        if sub_group:
            subtitle_groups.append(list(sub_group))
        continue

    # Check if adding this sentence would exceed limits
    if current_cjk + sent_cjk > 13 or current_lines >= 2:
        if current_group:
            subtitle_groups.append(list(current_group))
        current_group = [sentence]
        current_cjk = sent_cjk
        current_lines = 1
    else:
        current_group.append(sentence)
        current_cjk += sent_cjk
        current_lines += 1

if current_group:
    subtitle_groups.append(list(current_group))

print(f"Subtitle groups: {len(subtitle_groups)}")
for i, g in enumerate(subtitle_groups):
    print(f"  G{i+1}: {' / '.join(g[:2])}")

# === Distribute timing across groups ===
# Use character ratio within original 9-segment structure
seg_structure = [0, 3, 6, 8, 11, 14, 17, 20, 23, 26]  # sentence indices per segment
# Structure: seg1: sentences 0-2, seg2: 3-5, seg3: 6-7, seg4: 8-10, seg5: 11-13, seg6: 14-16, seg7: 17-19, seg8: 20-22, seg9: 23-25

# Calculate group CJK weights for timing
group_cjk = [sum(total_cjk(l) for l in g) for g in subtitle_groups]
total_groups_cjk = sum(group_cjk)
group_boundaries = [0.0]
for i, gc in enumerate(group_cjk):
    group_boundaries.append(group_boundaries[-1] + (gc / total_groups_cjk) * TOTAL_DURATION)
group_boundaries[-1] = TOTAL_DURATION

# Assign group -> segment mapping
# Figure out which groups belong to which segment
group_to_seg = []
for gi, g in enumerate(subtitle_groups):
    group_to_seg.append(gi * 9 // len(subtitle_groups))  # approximate

# === Helper functions ===
def format_time(sec):
    h = int(sec // 3600)
    m = int((sec % 3600) // 60)
    s = sec % 60
    return f'{h:02d}:{m:02d}:{s:06.3f}'

def srt_time(sec):
    return format_time(sec).replace('.', ',')

def ass_time(sec):
    return format_time(sec)

def get_kw_color(text):
    for kw in sorted(KEYWORDS, key=len, reverse=True):
        if kw in text:
            return KW_COLORS.get(kw, '#FFD700')
    return None

def highlight_text(text):
    """Apply ASS color tags for keywords."""
    sorted_kws = sorted(KEYWORDS, key=len, reverse=True)
    result = []
    remaining = text
    while remaining:
        best_pos = len(remaining)
        best_kw = None
        for kw in sorted_kws:
            pos = remaining.find(kw)
            if pos != -1 and pos < best_pos:
                best_pos = pos
                best_kw = kw
        if best_kw is not None:
            if best_pos > 0:
                result.append(remaining[:best_pos])
            color = KW_COLORS.get(best_kw, '#FFD700')
            bgr = color[5:7] + color[3:5] + color[1:3]
            result.append(fr"{{\c&H{bgr}&\b1}}{best_kw}{{\c&HFFFFFF&\b0}}")
            remaining = remaining[best_pos + len(best_kw):]
        else:
            result.append(remaining)
            break
    return ''.join(result)

# === Generate SRT ===
srt_lines = []
for i, (group, g_start, g_end) in enumerate(zip(subtitle_groups, group_boundaries[:-1], group_boundaries[1:])):
    if g_end - g_start < 0.5:
        g_end = g_start + 0.5
    srt_lines.append(str(i+1))
    srt_lines.append(f'{srt_time(g_start)} --> {srt_time(g_end)}')
    for line in group:
        srt_lines.append(line)
    srt_lines.append('')

srt_path = os.path.join(SUB_DIR, 'captions_final_v2.srt')
with open(srt_path, 'w', encoding='utf-8') as f:
    f.write('\n'.join(srt_lines))
print(f'\nSRT saved: {len(subtitle_groups)} entries')

# === Generate ASS ===
ass_header = '''[Script Info]
ScriptType: v4.00+
PlayResX: 1080
PlayResY: 1920
ScaledBorderAndShadow: yes

[V4+ Styles]
Format: Name, Fontname, Fontsize, PrimaryColour, SecondaryColour, OutlineColour, BackColour, Bold, Italic, Underline, StrikeOut, ScaleX, ScaleY, Spacing, Angle, BorderStyle, Outline, Shadow, Alignment, MarginL, MarginR, MarginV, Encoding
Style: Default,Microsoft YaHei,30,&H00FFFFFF,&H0000FF00,&H00000000,&H80000000,0,0,0,0,100,100,0,0,1,2,0,2,50,50,80,134
Style: Keyword,Microsoft YaHei,30,&H0000FFDD,&H0000FF00,&H00000000,&H80000000,1,0,0,0,100,100,0,0,1,2,0,2,50,50,80,134

[Events]
Format: Layer, Start, End, Style, Name, MarginL, MarginR, MarginV, Effect, Text
'''

ass_lines = [ass_header]
for i, (group, g_start, g_end) in enumerate(zip(subtitle_groups, group_boundaries[:-1], group_boundaries[1:])):
    if g_end - g_start < 0.5:
        g_end = g_start + 0.5
    # Check if any line has highlight
    has_hl = any(get_kw_color(l) for l in group)
    style_name = 'Keyword' if has_hl else 'Default'
    hl_lines = [highlight_text(l) for l in group]
    ass_text = '\\N'.join(hl_lines)
    ass_lines.append(f'Dialogue: 0,{ass_time(g_start)},{ass_time(g_end)},{style_name},,0,0,0,,{ass_text}')

ass_path = os.path.join(SUB_DIR, 'captions_final_v2.ass')
with open(ass_path, 'w', encoding='utf-8') as f:
    f.write('\n'.join(ass_lines))
print(f'ASS saved: {len(subtitle_groups)} events')

# === Generate word-level JSON ===
segments_json = []
current_time = 0.0
for group, g_start, g_end in zip(subtitle_groups, group_boundaries[:-1], group_boundaries[1:]):
    combined_text = ''.join(group)
    chars = list(combined_text)
    num_chars = max(len(chars), 1)
    char_dur = (g_end - g_start) / num_chars

    words = []
    for ci, ch in enumerate(chars):
        w_start = g_start + ci * char_dur
        w_end = g_start + (ci + 1) * char_dur
        words.append({
            'word': ch,
            'start': round(w_start, 2),
            'end': round(w_end, 2),
            'is_highlight': get_kw_color(ch) is not None
        })

    segments_json.append({
        'start': round(g_start, 2),
        'end': round(g_end, 2),
        'text': '\n'.join(group),
        'words': words
    })

captions_words = {
    'total_duration': round(TOTAL_DURATION, 1),
    'model': 'faster-whisper-timed + smart-grouping',
    'segments': segments_json,
    'word_count': sum(len(s['words']) for s in segments_json),
    'subtitle_groups': len(subtitle_groups),
    'keywords': KEYWORDS
}

words_path = os.path.join(SUB_DIR, 'captions_words_v2.json')
with open(words_path, 'w', encoding='utf-8') as f:
    json.dump(captions_words, f, ensure_ascii=False, indent=2)
print(f'captions_words_v2.json saved: {captions_words["word_count"]} words')

print('\n=== V2 SUBTITLE GENERATION COMPLETE ===')
print(f'Groups: {len(subtitle_groups)}, Duration: {TOTAL_DURATION}s')